# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：sourcesink.py
    Created on：2025/4/9 20:24
    Author：(input)
    Description:
"""
import numpy as np
import scipy.special as sp
import pandas as pd

# 源汇组合类（一源一汇）
class SourceSink:
    def __init__(self, sosi_strength, z_source, z_sink, C, k, length, height):
        """
        初始化一源一汇流动参数，源汇位置在物理平面上，通过 SC 变换映射到辅助平面

        Parameters
        ----------
        sosi_strength : tuple or list
            Strength of the source and sink pair (source_strength, sink_strength) in m^2/s
        z_source : complex
            Source position in the physical plane (x1 + i y1)
        z_sink : complex
            Sink position in the physical plane (x2 + i y2)
        C : float
            Parameter for the Schwarz-Christoffel (SC) transformation function
        k : float
            Modulus number for elliptic functions
        length : float
            Length of room (m)
        height : float
            Height of room (m)
        """
        # 参数验证
        if not isinstance(sosi_strength, (tuple, list)) or len(sosi_strength) != 2:
            raise ValueError("sosi_strength must be a tuple or list with 2 elements (source, sink)")
        if C is None or not isinstance(C, (int, float)):
            raise ValueError("C must be a valid number (int or float), not None")
        if k is None or not isinstance(k, (int, float)):
            raise ValueError("k must be a valid number (int or float), not None")
        if length <= 0:
            raise ValueError("房间长度 length 必须大于 0")
        if height <= 0:
            raise ValueError("房间高度 height 必须大于 0")
        if not isinstance(z_source, complex):
            raise ValueError("z_source must be a complex number (x1 + i y1)")
        if not isinstance(z_sink, complex):
            raise ValueError("z_sink must be a complex number (x2 + i y2)")

        # 源和汇的强度
        self.Q_source = sosi_strength[0]
        self.Q_sink = sosi_strength[1]

        # 房间尺寸
        self.L = length
        self.H = height

        # 源和汇在物理平面上的坐标
        self.z_so = z_source
        self.z_si = z_sink

        # 检查源汇坐标是否在房间范围内
        if not (-self.L/2 <= self.z_so.real <= self.L/2 and 0 <= self.z_so.imag <= self.H):
            raise ValueError("Source coordinates must be within room dimensions")
        if not (-self.L/2 <= self.z_si.real <= self.L/2 and 0 <= self.z_si.imag <= self.H):
            raise ValueError("Sink coordinates must be within room dimensions")

        # SC 变换参数
        self.C = float(C)
        self.k = float(k)
        self.k_prime = np.sqrt(1 - self.k**2)
        self.epsilon = 1e-10

        # 计算源和汇在辅助平面上的坐标
        self.w_so = self._SCT_g(self.z_so.real, self.z_so.imag)
        self.w_si = self._SCT_g(self.z_si.real, self.z_si.imag)

    def __repr__(self):
        return (f"SourceSink(source_strength={self.Q_source}, sink_strength={self.Q_sink}, "
                f"z_source={self.z_so}, z_sink={self.z_si}, "
                f"w_source={self.w_so}, w_sink={self.w_si}, "
                f"C={self.C}, k={self.k}, length={self.L}, height={self.H})")

    def _SCT_g(self, X, Y):
        """
        计算复数域的椭圆正弦函数 sn(z, k)
        Parameters:
           physics
            X : float or np.ndarray
                Real part
            Y : float or np.ndarray
                Imaginary part
        Returns:
            zeta : complex or np.ndarray
                Elliptic sine function result
        """
        use_X = X / self.C
        use_Y = Y / self.C
        sn_X_k, cn_X_k, dn_X_k, _ = sp.ellipj(use_X, self.k**2)
        sn_Y_kp, cn_Y_kp, dn_Y_kp, _ = sp.ellipj(use_Y, self.k_prime**2)
        num = sn_X_k * dn_Y_kp + 1j * sn_Y_kp * cn_Y_kp * cn_X_k * dn_X_k
        den = cn_Y_kp**2 + self.k**2 * sn_X_k**2 * sn_Y_kp**2
        return num / den

    def _SCT_g_prime(self, X, Y):
        """
        计算椭圆正弦函数 sn(z, k) 的导数 cn(z, k) * dn(z, k)
        Parameters:
            X : float or np.ndarray
                Real part
            Y : float or np.ndarray
                Imaginary part
        Returns:
            dsndZ : complex or np.ndarray
                Derivative of elliptic sine function
        """
        use_X = X / self.C
        use_Y = Y / self.C
        sn_X_k, cn_X_k, dn_X_k, _ = sp.ellipj(use_X, self.k**2)
        sn_Y_k, cn_Y_k, dn_Y_k, _ = sp.ellipj(use_Y, self.k**2)
        sn_X_kp, cn_X_kp, dn_X_kp, _ = sp.ellipj(use_X, self.k_prime**2)
        sn_Y_kp, cn_Y_kp, dn_Y_kp, _ = sp.ellipj(use_Y, self.k_prime**2)
        sn_iY_k = 1j * (sn_Y_kp / cn_Y_kp)
        cn_iY_k = 1 / cn_Y_kp
        dn_iY_k = dn_Y_kp / cn_Y_kp
        CN_num = sn_X_k * cn_X_k * dn_iY_k - sn_iY_k * cn_iY_k * dn_X_k
        CN_den = sn_X_k * cn_iY_k * dn_iY_k - sn_iY_k * cn_X_k * dn_X_k
        CN_Z_k = CN_num / CN_den
        DN_num = dn_X_k * cn_Y_kp * dn_Y_kp - 1j * self.k**2 * sn_X_k * cn_X_k * sn_Y_kp
        DN_den = cn_Y_kp**2 + self.k**2 * sn_X_k**2 * sn_Y_k**2
        DN_Z_k = DN_num / DN_den
        return CN_Z_k * DN_Z_k / self.C

    def cplx_W(self, X, Y):
        """
        计算复位势及其分量
        Parameters:
            X : np.ndarray
                Cartesian X coordinates
            Y : np.ndarray
                Cartesian Y coordinates
        Returns:
            w_ptn : np.ndarray
                Complex potential
            phi : np.ndarray
                Potential function
            psi : np.ndarray
                Stream function
        """
        Z = X + 1j * Y
        w_ptn = np.zeros_like(Z, dtype=np.complex128)
        phi = np.zeros_like(Z, dtype=np.float64)
        psi = np.zeros_like(Z, dtype=np.float64)

        zeta = self._SCT_g(X, Y)
        zeta1 = self.w_so
        zeta2 = self.w_si

        w_ptn = (self.Q_source / (2 * np.pi)) * np.log(zeta - zeta1 + self.epsilon) + \
                (self.Q_sink / (2 * np.pi)) * np.log(zeta - zeta2 + self.epsilon)
        phi = w_ptn.real
        psi = w_ptn.imag

        return w_ptn, phi, psi

    def cplx_velocity(self, X, Y):
        """
        计算复势 w 对复坐标 Z 的导数
        Parameters:
            X : np.ndarray
                Cartesian X coordinates
            Y : np.ndarray
                Cartesian Y coordinates
        Returns:
            magV : np.ndarray
                Velocity magnitude
            u_vel : np.ndarray
                Horizontal velocity component
            v_vel : np.ndarray
                Vertical velocity component
        """
        Z = X + 1j * Y
        magV = np.zeros_like(Z, dtype=np.float64)
        u_vel = np.zeros_like(Z, dtype=np.float64)
        v_vel = np.zeros_like(Z, dtype=np.float64)

        zeta = self._SCT_g(X, Y)
        zeta1 = self.w_so
        zeta2 = self.w_si
        dsndz = self._SCT_g_prime(X, Y)

        dwdZ = dsndz * ((self.Q_source / (2 * np.pi)) / (zeta - zeta1 + self.epsilon) +
                        (self.Q_sink / (2 * np.pi)) / (zeta - zeta2 + self.epsilon))

        magV = np.abs(dwdZ)
        u_vel = dwdZ.real
        v_vel = -dwdZ.imag

        return magV, u_vel, v_vel

    def pressure(self, U_ref, magV, u_vel, v_vel):
        """
        计算压力分布和压力系数 CP
        Parameters:
            U_ref : float
                Reference velocity
            magV : np.ndarray
                Velocity magnitude
            u_vel : np.ndarray
                Horizontal velocity component
            v_vel : np.ndarray
                Vertical velocity component
        Returns:
            p : np.ndarray
                Pressure distribution
            CP : np.ndarray
                Pressure coefficient distribution
            d_v2 : np.ndarray
                Velocity error term (for diagnostics)
        """
        if not (magV.shape == u_vel.shape == v_vel.shape):
            raise ValueError("Input arrays magV, u_vel, and v_vel must have the same shape")

        rho = 1.2  # 空气密度 (kg/m^3)
        V_V = u_vel ** 2 + v_vel ** 2
        d_v2 = magV ** 2 - V_V
        p = 0.5 * rho * (U_ref ** 2 - V_V)
        CP = p / (0.5 * rho * U_ref ** 2) - 0.99999

        return p, CP, d_v2

    def export_field_data(self, X, Y, U_ref=1.0, to_csv=False, filename="field_data.csv"):
        """
        输出每个点的坐标和对应的物理量：φ, ψ, magV, u, v, p, Cp 等

        Parameters:
            X : np.ndarray
                X 坐标网格（与 meshgrid 生成的 X 相同）
            Y : np.ndarray
                Y 坐标网格（与 meshgrid 生成的 Y 相同）
            U_ref : float
                参考速度（默认 1.0）
            to_csv : bool
                是否导出为 CSV 文件
            filename : str
                导出的 CSV 文件名（仅当 to_csv=True 时有效）

        Returns:
            df : pd.DataFrame
                包含所有结果的表格
        """
        import pandas as pd

        w, phi, psi = self.cplx_W(X, Y)
        magV, u_vel, v_vel = self.cplx_velocity(X, Y)
        p, CP, d_v2 = self.pressure(U_ref, magV, u_vel, v_vel)

        # 将数据整理成平铺形式
        data = {
            "x": X.flatten(),
            "y": Y.flatten(),
            "phi": phi.flatten(),
            "psi": psi.flatten(),
            "magV": magV.flatten(),
            "u": u_vel.flatten(),
            "v": v_vel.flatten(),
            "p": p.flatten(),
            "Cp": CP.flatten(),
            "velocity_error": d_v2.flatten()
        }

        df = pd.DataFrame(data)

        if to_csv:
            try:
                df.to_csv(filename, index=False)
                print(f"[INFO] 数据已导出到 {filename}")
            except PermissionError:
                print(f"[ERROR] 无法写入文件 '{filename}'，请确认该文件未被其他程序（如 Excel）打开，然后重试。")

        return df
